import random
import tkinter
from tkinter.filedialog import *
import json
import pygame
from pygame.key import *
from pygame.locals import *
from pygame.color import *

import pymunk
import pymunk.constraints
import pymunk.pygame_util


def build_envirment():  # 创建边框环境
    static_body = space.static_body
    static_lines = [
        pymunk.Segment(static_body, (0.0, 0.0), (0.0, 600.0), 3.0),
        pymunk.Segment(static_body, (0.0, 0.0), (1200, 0.0), 3.0),
        pymunk.Segment(static_body, (1200.0, 0.0), (1200.0, 600.0), 3.0),
        pymunk.Segment(static_body, (0.0, 600.0), (1200.0, 600.0), 3.0),

    ]

    for line in static_lines:
        line.elasticity = 0.95  # 弹性系数 [0, 1]
        line.friction = 0.4  # 摩擦系数 [0, 1]
        space.add(line)


def open_map():  # 打开已有的血管图
    root = tkinter.Tk()
    root.withdraw()
    save_dir = ''
    try:
        file_path = tkinter.filedialog.askopenfilename(title=u'选择文件', initialdir='.\\')
        json_file = open(file_path)
        draw_list = json.load(json_file)
        for key, i in enumerate(lines_in_space):  # 移除现有的所有线
            space.remove(i)
        for key, i in enumerate(balls):  # 移除现有的所有球
            space.remove(i)
        for key, i in enumerate(connectors):  # 移除现有的所有链接
            for j in i:
                space.remove(j)
        # 清空缓存列表
        lines_in_space.clear()
        balls.clear()
        connectors.clear()
        # 从缓存中绘制管路
        for i in draw_list:
            draw_a_line(i[0], i[1])
        # 将导丝重新绘入环境中
        create_guidewire(start_angle=start_ang)
    except:
        print("取消读取")


def save_map():  # 将当前血管图保存到文件
    root = tkinter.Tk()
    root.withdraw()
    fname = tkinter.filedialog.asksaveasfilename(title=u'保存文件', filetypes=[("json", ".json")], initialdir='.\\')
    save_list = []
    for i in lines_in_space:
        save_list.append([i.a.int_tuple, i.b.int_tuple])
    jsonArr = json.dumps(save_list, ensure_ascii=False)
    try:
        if fname == "":
            print("取消保存")
            return 0
        if fname[-5::] == '.json':
            fname = fname[::-5]
        with open(f'{fname}副本.json', 'w') as json_file:
            json_file.seek(0)
            json_file.truncate()
            json_file.write(jsonArr)
        print(f"已保存到{fname}副本.json")
    except:
        print(fname)
        pass


def draw_a_line(p1, p2):
    static_body = space.static_body
    line = pymunk.Segment(static_body, p1, p2, 5.0)
    line.elasticity = 0.95  # 弹性系数 0-1
    line.friction = 0.4  # 摩擦系数 0-1
    # lines_points.append([p1, p2])
    lines_in_space.append(line)
    space.add(line)


# 创建一个动态球
def create_dynamic_ball(mass=1, radius=5, position=(300, 300),
                        elasticity=0.1, friction=0.3, insert=False):  # 是否插入到首位
    inertia = pymunk.moment_for_circle(mass, 0, radius, position)
    body = pymunk.Body(mass, inertia)
    body.position = position[0], position[1]
    shape = pymunk.Circle(body, radius, (0, 0))
    shape.elasticity = elasticity
    shape.friction = friction
    space.add(body, shape)
    if insert:
        balls.insert(1, shape)
    else:
        balls.append(shape)
    return body


# 创建一个静止球
def create_static_ball(mass=1, radius=5, position=(300, 300), elasticity=0.9, friction=0.5, ):
    inertia = pymunk.moment_for_circle(mass, 0, radius, position)
    body = space.static_body
    body.position = position[0], position[1]
    shape = pymunk.Circle(body, radius, (0, 0))
    shape.elasticity = elasticity
    shape.friction = friction
    space.add(shape)
    balls.append(shape)

    return body


# 创建一个可变球
def create_kinamec_ball(mass=1, radius=5, position=(300, 300), elasticity=0.9, friction=0.5, ):
    inertia = pymunk.moment_for_circle(mass, 0, radius, position)
    body = pymunk.Body(mass, inertia, body_type=pymunk.Body.KINEMATIC)
    body.position = position[0], position[1]
    shape = pymunk.Circle(body, radius, (0, 0))
    shape.elasticity = elasticity
    shape.friction = friction
    space.add(shape, body)
    if not balls:
        balls.append(shape, )
    else:
        balls.insert(0, shape)
    return body


def insert_a_ball(radius=7):  # 插入一个球
    # 删除上一个球
    ini_pos = balls[0].body.position
    space.remove(balls[0])
    balls.pop(0)
    for i in connectors[0]:
        space.remove(i)
    connectors.remove(connectors[0])
    o = create_kinamec_ball(radius=radius, position=ini_pos)  # 新的可变球,永远在列表0的位置
    b = create_sec(o, radius=radius, ela=1e8, dump=1e9, insert=True)  # 新的链节，在列表1的位置
    if len(balls) < 8:
        c = create_connect(b, balls[2].body, radius=radius, angle=start_ang, stright=90)
    else:
        c = create_connect(b, balls[2].body, radius=radius, angle=0, stright=90)
    connectors.insert(1, c)


def pull_back_a_ball():  # 撤回一个球
    if len(balls) > 2:
        space.remove(balls[1])
        balls.pop(1)
        for i in connectors[0]:
            space.remove(i)
        for i in connectors[1]:
            space.remove(i)
        connectors.pop(0)
        connectors.pop(0)
        c = create_connect(balls[0].body, balls[1].body, 0, 90)
        connectors.insert(0, c)


def create_connect(prev_body, next_body,
                   radius=5, angle=0, stright=30, ela=1e8, dump=1e7):  # 创造两个球之间的连接
    string = pymunk.constraints.DampedRotarySpring(prev_body,
                                                   next_body,
                                                   angle * 3.14 / 180, ela, dump)  # 弹性副
    pivot = pymunk.constraints.PivotJoint(prev_body, next_body, (radius+1, 0), (-radius-1, 0), )  # 旋转副
    pivot.collide_bodies = True
    joint = pymunk.constraints.RotaryLimitJoint(prev_body,
                                                next_body,
                                                -stright * 3.14 / 180, stright * 3.14 / 180)  # 限制副
    space.add(string)
    space.add(pivot)
    space.add(joint)
    return [string, pivot, joint]


def create_sec(prev_body, radius=5, angle=0, stright=90, ela=1e8, dump=1e7, insert=False):  # 创建一个链节
    b = create_dynamic_ball(radius=radius, position=(prev_body.position[0] + 2*radius+1, prev_body.position[1]),
                            insert=insert)
    c = create_connect(prev_body, b, radius=radius, angle=angle, stright=stright, ela=ela, dump=dump)
    if insert:
        connectors.insert(0, c)
    else:
        connectors.append(c)  # 存储连接方式
    return b


def create_guidewire(start_position=(5, 345), start_angle=-45):  # 创建一个初始导丝
    o = create_kinamec_ball(position=start_position)
    b = create_sec(o, angle=0, stright=90)

    for i in range(5):  # 创造一节直的
        b = create_sec(b, angle=0, stright=90)
    for i in range(5):  # 创造一节可弯的
        b = create_sec(b, angle=start_angle, stright=90)


def my_events():  # 键盘事件监听
    global moving
    global start_pos
    global end_pos
    global start_ang

    for event in pygame.event.get():
        if event.type == QUIT:
            exit()
        elif event.type == pygame.MOUSEBUTTONDOWN:  # 获取点击鼠标事件
            if event.button == 1:  # 点击鼠标左键
                moving = True
            if event.button == 3:  # 点击鼠标右键
                start_pos = pygame.mouse.get_pos()
            if event.button == 2:  # 点击鼠标中键
                obj = space.point_query_nearest(pygame.mouse.get_pos(), 0, pymunk.ShapeFilter())
                if obj and obj.shape in lines_in_space:  # 删除线条
                    lines_in_space.remove(obj.shape)
                    space.remove(obj.shape)

        elif event.type == pygame.MOUSEBUTTONUP:  # 获取松开鼠标事件
            if event.button == 1:  # 松开鼠标左键
                moving = False
            if event.button == 3:  # 点击鼠标右键
                end_pos = pygame.mouse.get_pos()
                draw_a_line(start_pos, end_pos)

        elif event.type == KEYDOWN and event.key == K_ESCAPE:
            exit()

        elif event.type == KEYDOWN and event.key == K_o:
            if len(balls) > 6 and start_ang < 45:
                for i in connectors[-5:]:
                    i[0].rest_angle += 10 * 3.14 / 180
                connectors[-6][0].rest_angle = 0.5 * connectors[-5][0].rest_angle
                start_ang += 10
                print(start_ang)
        elif event.type == KEYDOWN and event.key == K_p:
            if len(balls) > 6 and start_ang > -45:
                for i in connectors[-5:]:
                    i[0].rest_angle -= 10 * 3.14 / 180
                connectors[-6][0].rest_angle = 0.5 * connectors[-5][0].rest_angle
                start_ang -= 10
                print(start_ang)
        elif event.type == KEYDOWN and event.key == K_d:
            # balls[0].body.position += pymunk.vec2d.Vec2d(10, 0)
            if len(balls) <= 10:
                insert_a_ball(radius=5)
            else:
                insert_a_ball()
        elif event.type == KEYDOWN and event.key == K_a:
            # balls[0].body.position += pymunk.vec2d.Vec2d(-10, 0)
            pull_back_a_ball()
        elif event.type == KEYDOWN and event.key == K_w:
            balls[0].body.position += pymunk.vec2d.Vec2d(0, -10)
        elif event.type == KEYDOWN and event.key == K_s:
            balls[0].body.position += pymunk.vec2d.Vec2d(0, 10)
        
        elif event.type == KEYDOWN and event.key == K_r:
            print(pygame.mouse.get_pos())
            create_static_ball(radius=15, position=pygame.mouse.get_pos())
        elif event.type == KEYDOWN and event.key == K_F1:
            save_map()
        elif event.type == KEYDOWN and event.key == K_F2:
            open_map()


def coll_begin(arbiter, space, data):
    print("begin", arbiter.shapes)
    return True

# 清除屏幕pymunk
def clear_screen():
    screen.fill(THECOLORS["white"])


# 渲染对象
def draw_objects():
    space.debug_draw(draw_options)


def main_loop():
    while True:  # 主循环
        # 计算下一帧
        for x in range(exact):
            space.step(1 / FPS / exact)

        my_events()  # 键盘事件

        clear_screen()  # 清空屏幕
        screen.blit(bg, (0,0))
        draw_objects()  # 绘制新屏幕
        if moving:
            balls[0].body.position = pygame.mouse.get_pos()  # 更新圆心位置为鼠标当前位置
            print(pygame.mouse.get_pos())

        pygame.display.flip()  # 时序控制
        clock.tick(FPS)
        pygame.display.set_caption("fps: " + str(clock.get_fps()))


if __name__ == "__main__":
    # -----------------------------------初始化环境------------------------------
    # pymunk初始化
    space = pymunk.Space()  # 物理引擎空间
    handler = space.add_default_collision_handler()  # 碰撞检测
    handler.begin = coll_begin


    # pygame初始化
    pygame.init()
    screen = pygame.display.set_mode((700, 600))
    clock = pygame.time.Clock()
    # 在pygame中创建画板
    draw_options = pymunk.pygame_util.DrawOptions(screen)
    bg = pygame.image.load("img_rsc\\3.png")
    moving = False
    start_pos = 0, 0
    end_pos = 0, 0

    # ----------------------------------环境参数---------------------------------
    space.gravity = (0.0, 0.0)  # 设置重力
    FPS = 60
    balls = []  # 所有的球
    lines_in_space = []
    connectors = []  # 所有的连接
    exact = 10  # 一帧计算几次
    start_ang = -45

    # ---------------------------------程序主体----------------------------------
    build_envirment()
    create_guidewire(start_angle=start_ang)
    main_loop()
